# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2021/9/23 上午11:15
# @Author  : acedar
# @FileName: lr_multi_classifer.py
# ----------------------------

# **** 评估指标 *******************
# ********** "f1" (default)
# ********** weightedPrecision
# ********** weightedRecall
# ********** accuracy
# *********************************

from pyspark.sql import SparkSession
from pyspark.ml.pipeline import Pipeline
from pyspark.ml.classification import DecisionTreeClassificationModel, DecisionTreeClassifier
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
from pyspark.ml.feature import IndexToString, StringIndexer, VectorIndexer

spark = SparkSession.builder.appName("test").getOrCreate()
data_path = "../datasets/mllib/sample_libsvm_data.txt"
data_df = spark.read.format("libsvm").load(data_path)

label_indexer = StringIndexer().setInputCol("label").setOutputCol("indexedLabel")\
    .fit(dataset=data_df)

feature_index = VectorIndexer().setInputCol("features").setOutputCol("indexedFeatures")\
    .setMaxCategories(4).fit(dataset=data_df)


train_df, test_df = data_df.randomSplit([0.8, 0.2])

dtc = DecisionTreeClassifier().setLabelCol("indexedLabel").setFeaturesCol("indexedFeatures")

label_convert = IndexToString().setInputCol("prediction").setOutputCol("predictedLabel")\
    .setLabels(label_indexer.labels)

pipepline = Pipeline().setStages([label_indexer, feature_index, dtc, label_convert])

# train
model = pipepline.fit(train_df)

# predict
pred_df = model.transform(test_df)

pred_df.select("predictedLabel", "label", "features").show(5)

multi_evaluator = MulticlassClassificationEvaluator().setLabelCol("indexedLabel")\
    .setPredictionCol("prediction").setMetricName("accuracy")

accuracy = multi_evaluator.evaluate(pred_df)
print(f"mulit test accuracy: {accuracy}")

print(f"multi error = {1.0-accuracy}")
